{-# OPTIONS -fglasgow-exts #-}
{-# OPTIONS -XPatternSignatures #-}

module Data.Transform.RelTransfHUnit where

import HUnit
import Data.Transform.RulesPF
import Data.Transform.Rewriting (Rule, reduceIO)
import Data.Transform.Type
import Data.Set as Set
import Data.Map as Map


main = runTestTT . TestList . Prelude.map TestCase $
   [ testRelId_1
   , testRelId_2
   , testRelId_3
   , testRelId_4
   , testRelId_5
   , testRelId_6
   , testRelId_7
   , testProdElim_1
   , testProdElim_2
   , testIdElim_1
   , testIdElim_2
   , testAssocr
   , testAssocl
   , testSplitElim_1
   , testSplitElim_2
   , testSplitElim_3
   , testSplitSimpl_1
   , testSplitSimpl_2
   , testCompDistr_1
   , testCompDistr_2
   , testCompDistr_3
   , testCompDistr_4
   , testCompDistr_5
   , testCompDistr_6
--   , testCompDistr_7 -- Failing!!
   , testCompDistr_8
--   , testCompDistr_9 -- Failing!!
--   , testCompDistr_10 -- Failing!!
   , testDomMapElim_1
   , testDomMapElim_2
   , testRngMapElim_1
   , testRngMapElim_2
   ]
   

testRelId_1 = do
   let t        = Func (Set Int) (Set Int)
   let actual   = SET ID
   let expected = ID
   checkRuleWithValues setmap_id t actual expected [Set.fromList [1,2,3]] 

testRelId_2 = do
   let t        = Func (List Int) (List Int)
   let actual   = LIST ID
   let expected = ID
   let values   = [[1,2,3], [1,2], [0,99,1000]]
   checkRuleWithValues listmap_id t actual expected values

testRelId_3 = do   
   let t        = Func (Maybe Int) (Maybe Int)
   let actual   = MAYBE ID
   let expected = ID
   checkRuleWithValues maybe_id t actual expected [Nothing, Just 2, Just 99]

testRelId_4 = do
   let t        = Func (Map Int Int) (Map Int Int)
   let actual   = MAPKEYS ID
   let expected = ID
   let values   = [ Map.fromList [(1,1), (2,1)], Map.fromList [(2,2), (3,4)]]
   checkRuleWithValues mapkeys_id t actual expected values

testRelId_5 = do
   let t        = Func (Map Int Int) (Map Int Int)
   let actual   = MAP ID
   let expected = ID
   let values   = [ Map.fromList [(1,4), (2,6)], Map.fromList [(0,0)]]
   checkRuleWithValues mapmap_id t actual expected values

testRelId_6 = do
   let t        = Func (Either Int String) (Either Int String)
   let actual   = SUM ID ID
   let expected = ID
   let values   = [Left 3, Right "4", Left 0, Right "hello"]
   checkRuleWithValues either_id t actual expected values
   
testRelId_7 = do
   let t        = Func (Prod String Int) (Prod String Int) 
   let actual   = PROD ID ID
   let expected = ID
   let values   = [("2",1), ("hello",0), ("9", 10)]
   checkRuleWithValues prod_id t actual expected values

testProdElim_1 = do
   let t        = Func (Prod Int Int) String
   let f        = FUN "show" (\ x -> show x)
   let actual   = COMP (Prod String Int) FST (PROD f ID)
   let expected = COMP Int f FST
   let values   = [(1,1), (100, 9), (13, 0)]
   checkRuleWithValues prod_elim t actual expected values

testProdElim_2 = do
   let t        = Func (Prod Int Int) String
   let f        = FUN "show" (\ x -> show x)
   let actual   = COMP (Prod Int String) SND (PROD ID f)
   let expected = COMP Int f SND
   let values   = [(1,1), (100, 9), (13, 0)]
   checkRuleWithValues prod_elim t actual expected values

testIdElim_1 = do
   let t        = Func (Prod String Int) String
   let actual   = COMP String ID FST
   let expected = FST
   let values   = [("a",1), ("b",2),("c",4)]
   checkRuleWithValues comp_id t actual expected values

testIdElim_2 = do
   let t        = Func (Prod String Int) Int
   let actual   = COMP (Prod String Int) SND ID
   let expected = SND
   let values   = [("a",1), ("b",2),("c",4)]
   checkRuleWithValues comp_id t actual expected values
   
testAssocr = do
   let t        = Func (Prod String (Prod Int Int)) Int
   let actual   = COMP (Prod Int Int) (COMP Int ID FST) SND
   let expected = COMP Int ID (COMP (Prod Int Int) FST SND)
   let values   = [("", (1,0)), ("test", (2,2))]
   checkRuleWithValues comp_assocr t actual expected values

testAssocl = do
   let t         = Func (Prod String (Prod Int Int)) Int
   let actual    = COMP Int ID (COMP (Prod Int Int) FST SND)
   let expected  = COMP (Prod Int Int) (COMP Int ID FST) SND 
   let values    = [("", (1,0)), ("test", (2,2))] 
   checkRuleWithValues comp_assocl t actual expected values

testSplitElim_1 = do
   let t        = Func Int Int
   let actual   = COMP (Prod Int One) FST (SPLIT ID BANG)
   let expected = ID
   let values   = [1,2,4,99]
   checkRuleWithValues prod_cancel t actual expected values

testSplitElim_2 = do
   let t        = Func Int Int
   let actual   = COMP (Prod One Int) SND (SPLIT BANG ID)
   let expected = ID
   let values   = [1,2,4,99]
   checkRuleWithValues prod_cancel t actual expected values
   
testSplitElim_3 = do
   let t        = Func (Prod Int String) (Prod Int String)
   let actual   = SPLIT FST SND
   let expected = ID
   let values   = [(1,"1.00"), (2,"20"), (9,"check")]
   checkRuleWithValues prod_reflex t actual expected values

testSplitSimpl_1 = do
   let t        = Func (Prod Int String) (Prod Int Int)
   let actual   = SPLIT (COMP Int ID FST) (COMP Int ID FST)
   let expected = COMP Int (SPLIT ID ID) FST
   let values   = [(2,"test"), (99,"0.33"), (1000,"!as!")]
   checkRuleWithValues split_simpl t actual expected values

testSplitSimpl_2 = do
   let t        = Func (Prod Int String) (Prod String String)
   let actual   = SPLIT (COMP String ID SND) (COMP String ID SND)
   let expected = COMP String (SPLIT ID ID) SND
   let values   = [(2,"test"), (99,"0.33"), (1000,"!as!")]
   checkRuleWithValues split_simpl t actual expected values

testCompDistr_1 = do
   let t        = Func (Set (Prod (Prod String Int) Int)) (Set String)
   let actual   = COMP (Set (Prod String Int)) (SET FST) (SET FST) 
   let expected = SET (COMP (Prod String Int) FST FST)
   let values   = [ Set.fromList [(("a", 0), 1), (("b", 1), 2)]]
   checkRuleWithValues comp_distr t actual expected values

testCompDistr_2 = do
   let t        = Func (List (Prod (Prod String Int) Int)) (List String)
   let actual   = COMP (List (Prod String Int)) (LIST FST) (LIST FST)
   let expected = LIST (COMP (Prod String Int) FST FST)
   let values   = [ [(("a", 0), 1), (("b", 1), 2)], [(("c", 2), 3)]]
   checkRuleWithValues comp_distr t actual expected values

testCompDistr_3 = do
   let t        = Func (Maybe (Prod (Prod String Int) Int)) (Maybe String)
   let actual   = COMP (Maybe (Prod String Int)) (MAYBE FST) (MAYBE FST)
   let expected = MAYBE (COMP (Prod String Int) FST FST)
   let values   = [Nothing, Just (("a", 0), 3), Just (("a",1), 9)]
   checkRuleWithValues comp_distr t actual expected values

testCompDistr_4 = do
   let t        = Func (List (Prod Int String)) (List Int)
   let actual   = COMP (List Int) (LIST ID) (LIST FST)
   let expected = LIST (COMP Int ID FST)
   let values   = [[(1,"a")], [(2,"b"), (3,"d")]]
   checkRuleWithValues comp_distr t actual expected values

testCompDistr_5 = do
   let t        = Func (Map (Prod Int Int) String) (Map Int String)
   let actual   = COMP (Map (Prod Int Int) String) (MAPKEYS FST) (MAPKEYS ID)
   let expected = MAPKEYS (COMP (Prod Int Int) FST ID)
   let values   = [Map.fromList [((1,1),"a"), ((2,3),"b")], Map.fromList [((4,4), "c")]]
   checkRuleWithValues comp_distr t actual expected values

testCompDistr_6 = do
   let t        = Func (Map Int (Prod String Int)) (Map Int String)
   let actual   = COMP (Map Int String) (MAP ID) (MAP FST)
   let expected = MAP (COMP String ID FST)
   let values   = [Map.fromList [(1,("a",3)), (3,("b",99))]]
   checkRuleWithValues comp_distr t actual expected values

testCompDistr_7 = do
   let t        = Func (Either (Prod Int Int) (Prod String String)) (Either Int String)
   let actual   = COMP (Either Int (Prod String String)) (SUM ID FST) (SUM FST ID)
   let expected = SUM (COMP Int ID FST) (COMP (Prod String String) FST ID)
   let values   = [Left (1,1), Right ("a","b")]
   checkRuleWithValues comp_distr t actual expected values

testCompDistr_8 = do
   let t        = Func (Prod Int Int) (Prod Int Int)
   let actual   = COMP (Prod Int Int) (PROD ID ID) (PROD ID ID)
   let expected = PROD (COMP Int ID ID) (COMP Int ID ID)
   let values   = [(1,1),(1,0),(100,900001)]
   checkRuleWithValues comp_distr t actual expected values

testCompDistr_9 = do
   let t        = Func (List (Prod Int String)) Bool
   let f        = FUN "Set.fromList" Set.fromList
   let actual   = COMP (List Int) (SetInc Int f f) (LIST FST)
   let expected = SetInc Int (COMP (List Int) f (LIST FST)) (COMP (List Int) f (LIST FST))
   let values   = [[(2,"a"),(3,"as")], [(1,"aaa")]]
   checkRuleWithValues comp_distr t actual expected values

testCompDistr_10 = do
   let t        = Func (Prod Int Int) Bool
   let f1       = FUN "GT 5" (\ x -> x > 5)
   let f2       = FUN "LT 100" (\ x -> x < 100)
   let actual   = COMP Int (AND f1 f2) FST
   let expected = AND (COMP Int f1 FST) (COMP Int f2 FST)
   let values   = [(1,1),(10,16),(9,0)]
   checkRuleWithValues comp_distr t actual expected values

testDomMapElim_1 = do
   let t        = Func (Map (Prod Int String) Int) (Set Int)
   let actual   = COMP (Map Int Int) DOM (MAPKEYS FST)
   let expected = COMP (Set (Prod Int String)) (SET FST) DOM
   let values   = [Map.fromList [((1, "a"), 1), ((2, "b"), 2)], Map.fromList [((3, "c"), 2)]]
   checkRuleWithValues mapkeys_fusion t actual expected values   

testDomMapElim_2 = do
   let t        = Func (Map (Prod Int Int) Int) (Set Int)
   let actual   = COMP (Map Int Int) RNG (MAPKEYS FST)
   let expected = RNG
   let values   = [Map.fromList [((1,1),2), ((2,1),3)], Map.fromList [((9,9),9)]]
   checkRuleWithValues mapkeys_fusion t actual expected values

testRngMapElim_1 = do
   let t        = Func (Map Int (Prod Int Int)) (Set Int)
   let actual   = COMP (Map Int Int) DOM (MAP FST)
   let expected = DOM
   let values   = [ Map.fromList [(1, (1,1)), (2, (3,4))], Map.fromList [(0,(0,0))]]
   checkRuleWithValues map_fusion t actual expected values

testRngMapElim_2 = do
   let t        = Func (Map Int (Prod Int String)) (Set String)
   let actual   = COMP (Map Int String) RNG (MAP SND)
   let expected = COMP (Set (Prod Int String)) (SET SND) RNG
   let values   = [ Map.fromList [(1, (1,"z")), (2, (3,"a"))], Map.fromList [(0,(0,"c"))]]
   checkRuleWithValues map_fusion t actual expected values
   
-- | Check a transformation rule by checking values.
checkRuleWithValues :: (Show b, Eq b) => Rule -> (Type (a -> b)) -> (PF (a -> b)) -> (PF (a -> b)) -> [a] -> Assertion
checkRuleWithValues rule t actual expected values = do
   result <- reduceIO rule t actual
   mapM_ (checkValue expected result) values


-- | Check if two functions provide the same result for the same value.
checkValue :: (Show b, Eq b) => (PF (a -> b)) -> (PF (a -> b)) -> a -> Assertion
checkValue originalPF resultPF value = do
    let expectedResult = eval originalPF value
    let actualResult   = eval resultPF value
    assertEqual "Values are not equal" expectedResult actualResult
     